from pylab import *
import numpy as np
import pandas as pd
import scipy.linalg as linalg
import matplotlib.pyplot as plt
from matplotlib.font_manager import FontProperties
fontP = FontProperties()
fontP.set_size('small')
from scikits.statsmodels.tools.tools import categorical
from scipy.stats import zscore

# our snow imports
from snow_similarity import get_similarity2index
from getXmostFilledRows import *

#-------------- READING THE DATA -----------------#

censusdata = np.loadtxt('./census-income/census-income.txt', dtype=np.str,delimiter=',')
censusdata = censusdata[:,:40]
attributeNames = np.loadtxt('./census-income/census_attributes.txt', dtype=np.str,delimiter=',')

#-------------- SPLITING THE CATEGORICAL DATA -----------------#

#Data that has to be split in to arrays:
selectedColumns = [0, 4, 5, 8, 9, 12, 14, 16, 17, 18, 20, 24, 30, 31, 34, 39]
censusdataSelected = censusdata[:, selectedColumns]
attributeNamesSelected = attributeNames[:, selectedColumns]

#Indices of columns that have to bee split into numerical
splitColumns = [1, 3, 4, 5, 6, 10, 12, 13, 14]
#Indices of columns that are continous
continousColumns = [0, 2, 7, 8, 9, 11, 15]

#Parameters for whole data
dataRows, dataColumns = np.shape(censusdata)
dataMatrix = censusdata.astype(str)

#Parameters for selectedData
dataRowsSel, dataColumnsSel = np.shape(censusdataSelected)
dataMatrixStrSel = censusdataSelected.astype(str)

##CATEGORICAL VALUES TO BINARY
#Splits nominal data into numerical values
n_of_nonSplitColumns = len(selectedColumns) - len(splitColumns)
class_idx = n_of_nonSplitColumns-1+1 # indicates where a binarized category starts from
class2index = {} # dictionary from classes to indexes at which they start in the new data matrix
for i in splitColumns:
	a = censusdataSelected[:,i]
        current_class = attributeNamesSelected[i].strip()
	b, dictNames = categorical(a, dictnames = True, drop=True)
        class2index[current_class] = (class_idx, len(dictNames))
        spanBeforeStacking = censusdataSelected.shape[1]
	censusdataSelected = np.hstack((censusdataSelected,b))

	for j in dictNames.keys():
		if dictNames[j].find("Not in") == -1:
	 		attributeNamesSelected = np.append(attributeNamesSelected, dictNames[j])
	 	else:     
	 		censusdataSelected = np.delete(censusdataSelected, spanBeforeStacking+j, axis=1)
            	class_idx -= 1
            	class2index[current_class] = (class2index[current_class][0],class2index[current_class][1]-1)
        
        class_idx += len(dictNames)

#Gets rid of the old columns
censusdataSelected = np.delete(censusdataSelected, splitColumns, axis=1)
attributeNamesSelected = np.delete(attributeNamesSelected, splitColumns)


#-------------- DATA REDUCTION -----------------#

# CONVERT THE MATRIX TO NUMERICAL
dataMatrix = censusdataSelected.astype(float)
instanceWeight_col=5

# EDUCATION ATTRIBUTES
basicSchoolColumns = [7,8,9,10,11,12,13,20]
bachelorOrUnderColumns = [14,15,16,19,23]
doctorColumns = [18, 22]
master = [21]
child = [17]
#Just for cleaning up the data
educationColumns = [7,8,9,10,11,12,13,14,15,16,18,19,20,22,23]

# INDUSTRY CODE
agricultureColumns = [24,32]
publicAndStateColumns = [25,29,41,43]
industryColumns = [26,27,34,35,37]
constructionColumns = [28,40,45]
financialAndConsultingColumns = [31,38,39]
medicalColumns = [33,36]
tradeAndLogisticsColumns = [42,43,46]
#Just for cleaning  up the data
industryCodeColumns = [24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46]

# OCCUPATION CODE
publicColumns = [47,48]
othersColumns = [50,53]
workerColumns = [51,52,54,55,60]
expertColumns = [56,59]
executiveColumns = [49]
salesColumns = [58]
#Just for cleaning up the data
occupationCodeColumns = [47,48,50,51,52,53,54,55,56,57,59,60]

#Other columns to delete
movedFromColumns = np.arange(63,73)
countryColumns = np.arange(80,127)
numberOfColumns = np.shape(attributeNamesSelected)
incomeColumns = [1,2,3,4,6]

#Combine the income columns to one
#def combineRowsToOne(dataToCombine, nbrOfWorkWeeks, salaryInHour, dividentsFromStocsk, capitalGains, capitalLosses):
combinedIncomeColumn = combineRowsToOne(dataMatrix[:,incomeColumns], 4, 0, 3, 1, 2)
dataMatrix = np.hstack((dataMatrix, combinedIncomeColumn))
attributeNamesSelected = np.append(attributeNamesSelected, 'Yearly income')

continousColumnsNew = [0,5,6,127]

#Returns the reduced data matrix wich has the most filled rows
#reducedDataMatrix = getXmostFilledRows(dataMatrix, 20000)
#By continous columns
reducedDataMatrix = getXmostFilledRowsByCol(dataMatrix, censusdata.shape[0], continousColumns)

#np.savetxt('reducedDataMatrix', reducedDataMatrix.astype(str), delimiter=',', fmt="%s")
reducedDataMatrix = combineCategoricals(reducedDataMatrix, basicSchoolColumns)
attributeNamesSelected = np.append(attributeNamesSelected, 'Edu. Basic School')
reducedDataMatrix = combineCategoricals(reducedDataMatrix, bachelorOrUnderColumns)
attributeNamesSelected = np.append(attributeNamesSelected, 'Edu. Bachelor or under')
reducedDataMatrix = combineCategoricals(reducedDataMatrix, doctorColumns)
attributeNamesSelected = np.append(attributeNamesSelected, 'Edu. Doctoral')
attributeNamesSelected[master] = 'Edu. Masters degree'
attributeNamesSelected[child] = 'Edu. Child'

reducedDataMatrix = combineCategoricals(reducedDataMatrix, agricultureColumns)
attributeNamesSelected = np.append(attributeNamesSelected, 'Inds. Agriculture')
reducedDataMatrix = combineCategoricals(reducedDataMatrix, publicAndStateColumns)
attributeNamesSelected = np.append(attributeNamesSelected, 'Inds. Public and state')
reducedDataMatrix = combineCategoricals(reducedDataMatrix, industryColumns)
attributeNamesSelected = np.append(attributeNamesSelected, 'Inds. Industry')
reducedDataMatrix = combineCategoricals(reducedDataMatrix, constructionColumns)
attributeNamesSelected = np.append(attributeNamesSelected, 'Inds. Construction')
reducedDataMatrix = combineCategoricals(reducedDataMatrix, financialAndConsultingColumns)
attributeNamesSelected = np.append(attributeNamesSelected, 'Inds. Financial and consulting')
reducedDataMatrix = combineCategoricals(reducedDataMatrix, medicalColumns)
attributeNamesSelected = np.append(attributeNamesSelected, 'Inds. Medical')
reducedDataMatrix = combineCategoricals(reducedDataMatrix, tradeAndLogisticsColumns)
attributeNamesSelected = np.append(attributeNamesSelected, 'Inds. Trade and logistics')

reducedDataMatrix = combineCategoricals(reducedDataMatrix, publicColumns)
attributeNamesSelected = np.append(attributeNamesSelected, 'Occu. Public worker')
reducedDataMatrix = combineCategoricals(reducedDataMatrix, othersColumns)
attributeNamesSelected = np.append(attributeNamesSelected, 'Occu. Other')
reducedDataMatrix = combineCategoricals(reducedDataMatrix, workerColumns)
attributeNamesSelected = np.append(attributeNamesSelected, 'Occu. Worker')
reducedDataMatrix = combineCategoricals(reducedDataMatrix, expertColumns)
attributeNamesSelected = np.append(attributeNamesSelected, 'Occu. Expert')
attributeNamesSelected[executiveColumns] = 'Occu. Executive degree'
attributeNamesSelected[salesColumns] = 'Occu. Child'

columnsToDelete = np.hstack((educationColumns,industryCodeColumns,occupationCodeColumns,movedFromColumns,countryColumns))
reducedDataMatrix = np.delete(reducedDataMatrix, columnsToDelete, axis=1)
attributeNamesSelected = np.delete(attributeNamesSelected, columnsToDelete)

#Enable if you want to save new data files
np.savetxt('census-income/reducedDataMatrix_03_full', reducedDataMatrix.astype(float), delimiter=',', fmt="%s")
np.savetxt('census-income/attributeNamesReduced_03_full', attributeNamesSelected.astype(str), delimiter=',', fmt="%s")

#-------------- STATISTICAL INFORMATION -----------------#

# CONVERT THE MATRIX TO NUMERICAL
dataMatrix = reducedDataMatrix.astype(str)
instanceWeight_col=5

numOfDifferentValuesStr = getNumberOfDifferentValues(dataMatrix, attributeNamesSelected)
print numOfDifferentValuesStr
print ""
